#!/usr/bin/env python

import os, sys
from astropy.io import fits
import numpy as np
from scipy.signal import medfilt2d
import glob
from astropy.modeling import models, fitting

def polyfit(data, degree=3):
    ny, nx = data.shape
#    yy, xx = np.mgrid[:ny, :nx]
    yy, xx = np.mgrid[-ny//2:ny//2, -nx//2:nx//2]
    p_init = models.Polynomial2D(degree=degree)
    fit_p = fitting.LinearLSQFitter()
#    fit_p = fitting.LevMarLSQFitter()
    p = fit_p(p_init, xx, yy, data)
    model = p(xx, yy)
    return model

def remove_overscan(imgdata):
#    img = np.zeros((9232,9216), dtype=np.uint16)
#    for i in range(4):
#        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+27):((i+1)*1250-71)]
#        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+27):((i+1)*1250-71)]
#    for i in range(4,8):
#        img[0:4616, i*1152:(i+1)*1152] = imgdata[0:4616,(i*1250+71):((i+1)*1250-27)]
#        img[4616:, i*1152:(i+1)*1152] = imgdata[4784:,(i*1250+71):((i+1)*1250-27)]
#    return img

    img = np.zeros((9230,9216), dtype=np.uint16)
    for i in range(4):
        img[0:4615, i*1152:(i+1)*1152] = imgdata[0:4615,(i*1250+27):((i+1)*1250-71)]
        img[4615:, i*1152:(i+1)*1152] = imgdata[4785:,(i*1250+27):((i+1)*1250-71)]
    for i in range(4,8):
        img[0:4615, i*1152:(i+1)*1152] = imgdata[0:4615,(i*1250+71):((i+1)*1250-27)]
        img[4615:, i*1152:(i+1)*1152] = imgdata[4785:,(i*1250+71):((i+1)*1250-27)]
    return img


def get_stacked_img(flist):
    fdata = None
    fdata_cnt = 0
    for f in flist:
        d = fits.getdata(f).astype(float)
        print('loaded : {}'.format(f))
        if fdata_cnt == 0:
            fdata = d
        else:
            fdata += d
        fdata_cnt += 1
    return fdata/fdata_cnt

def avg_filter(img,kernel_size=15):
    tmp = medfilt2d(img, kernel_size=kernel_size)
    index = tmp < 10
    tmp[index] = np.median(tmp)
    return tmp

def count_darkpixels(img):
#    img_filt = polyfit(img, degree=degree)
    img_filt = avg_filter(img)
    
    hdu = fits.PrimaryHDU(img_filt)
    hdu.writeto('flt_2d.fits', overwrite=True)
    
    ratio = img / img_filt
    
    hdu = fits.PrimaryHDU(ratio)
    hdu.writeto('ratio.fits', overwrite=True)
    
    ndark = np.sum(ratio<0.5) 
    frac = ndark / img.shape[0] / img.shape[1]
    print('ndark: {}'.format(ndark))
    print('frac: {}'.format(frac))

def main():

    flist = glob.glob('*paper*.fits')
#    flist = ['flat_nopining_paper_550nm_27s_1.fits',
#             'flat_nopining_paper_550nm_27s_2.fits',
#             'flat_nopining_paper_550nm_27s_3.fits',
#             'flat_nopining_paper_550nm_27s_4.fits',
#             'flat_nopining_paper_550nm_27s_5.fits']

    flat_stacked = get_stacked_img(flist)
    bias = fits.getdata('bias-10.fits')
    flat_stacked = flat_stacked - bias
    flat_stacked_ored = remove_overscan(flat_stacked)
    count_darkpixels(flat_stacked_ored)

if __name__ == '__main__':
    main()
